/**
 * Copyright 2019-2020 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "graph_optimizer/ub_fusion/tbe_pass/tbe_conv2d_backprop_elemwise_fusion_pass.h"
#include <algorithm>
#include <string>
#include <vector>

namespace fe {
using std::vector;
static const char PATTERN_DX[] = "conv2dBackpropInput";
static const char PATTERN_ELTWISE[] = "eltwise";
static const char PATTERN_OTHER_INPUT[] = "otherInput";
static const char PATTERN_ELTWISE1[] = "eltwise1";
static const char PATTERN_OTHER_INPUT1[] = "otherInput1";

/*
 * @brief:  define dx and two inputs elemwise op fusion pattern
 *
 *   Pattern:
 *   conv2d_backprop_input-->elemwise1-->elemwise2-->OpTypeAny
 *              other_input1--/          /
 *                        other_input2--/
 * fusion node: conv2d_backprop_input, elemwise1, elemwise2
 *
 *   Pattern1:
 *   conv2d_backprop_input-->elemwise-->OpTypeAny
 *                other_input--/
 * fusion node: conv2d_backprop_input, elemwise
 *
 * @return BufferFusionPattern: return all valid patterns.
 */
vector<BufferFusionPattern *> TbeConv2DBackpropElemwiseFusionPass::DefinePatterns() {
  vector<BufferFusionPattern *> patterns;

  string pass_name = "TbeConv2dBackpropInputTwoElemwiseFusion";
  BufferFusionPattern *pattern = new (std::nothrow) BufferFusionPattern(pass_name);
  FE_CHECK((pattern == nullptr),
            REPORT_FE_ERROR("[SubGraphOpt][Conv2DBkpropFus][DefPtn] %s new an object failed.", pass_name.c_str()),
            return patterns);

  // define pattern rules
  pattern->AddOpDesc(PATTERN_DX, {OP_PATTERN_CONV_BACKPROP_INPUT}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .AddOpDesc(PATTERN_ELTWISE, {OP_PATTERN_ELEMWISE}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .AddOpDesc(PATTERN_OTHER_INPUT, {TBE_PATTERN_INPUT_NODE}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .AddOpDesc(PATTERN_ELTWISE1, {OP_PATTERN_ELEMWISE}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .AddOpDesc(PATTERN_OTHER_INPUT1, {TBE_PATTERN_INPUT_NODE}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .SetHead({PATTERN_DX})
      .SetOutputs(PATTERN_DX, {PATTERN_ELTWISE})
      .SetOutputs(PATTERN_OTHER_INPUT, {PATTERN_ELTWISE})
      .SetOutputs(PATTERN_ELTWISE, {PATTERN_ELTWISE1})
      .SetOutputs(PATTERN_OTHER_INPUT1, {PATTERN_ELTWISE1});

  patterns.push_back(pattern);
  FE_LOGD("End to define %s pass pattern.", pass_name.c_str());

  string pass_name1 = "TbeConv2dBackpropInputElemwiseFusion";
  FE_LOGD("Start to define %s pass pattern.", pass_name1.c_str());
  BufferFusionPattern *pattern1 = new (std::nothrow) BufferFusionPattern(pass_name1);
  FE_CHECK((pattern1 == nullptr),
            REPORT_FE_ERROR("[SubGraphOpt][Conv2DBkpropFus][DefPtn] %s new an object failed.", pass_name1.c_str()),
            return patterns);

  // define pattern1 rules
  pattern1->AddOpDesc(PATTERN_DX, {OP_PATTERN_CONV_BACKPROP_INPUT}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .AddOpDesc(PATTERN_ELTWISE, {OP_PATTERN_ELEMWISE}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .AddOpDesc(PATTERN_OTHER_INPUT, {TBE_PATTERN_INPUT_NODE}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .SetHead({PATTERN_DX})
      .SetOutputs(PATTERN_DX, {PATTERN_ELTWISE})
      .SetOutputs(PATTERN_OTHER_INPUT, {PATTERN_ELTWISE});

  patterns.push_back(pattern1);
  FE_LOGD("End to define %s pass pattern.", pass_name1.c_str());

  return patterns;
}

/*
 * @brief: parse nodes matched in mapping and call DoFusion
 * @param [in] graph: original graph
 * @param [out] mapping: nodes matched by pattern
 * @return bool: fusion status ok or not.
 */
Status TbeConv2DBackpropElemwiseFusionPass::GetFusionNodes(const BufferFusionMapping &mapping,
                                                           vector<ge::NodePtr> &fusion_nodes) {
  FE_LOGD("Begin to do Conv2dBackpropInputElemwiseFusion!");
  vector<ge::NodePtr> elem1_node = GetMatchedNodesByDescName(PATTERN_ELTWISE1, mapping);
  vector<ge::NodePtr> elem_node = GetMatchedNodesByDescName(PATTERN_ELTWISE, mapping);
  if (elem_node.empty()) {
    FE_LOGW("Elemwise node not matched.");
    return SUCCESS;
  }
  if (!elem1_node.empty()) {
    bool check_elemwise = elem1_node[0]->GetType() == "ReluGradV2" && elem_node[0]->GetType() == "AddN" &&
                          (elem_node[0]->GetOpDesc()->GetInputDesc(0).GetDataType() ==
                           elem_node[0]->GetOpDesc()->GetInputDesc(1).GetDataType()) &&
                          (elem_node[0]->GetOpDesc()->GetInputDesc(0).GetShape().GetDims() ==
                           elem_node[0]->GetOpDesc()->GetInputDesc(1).GetShape().GetDims());
    if (!check_elemwise) {
      FE_LOGD("The optype of node[%s] and [%s] should be AddN and ReluGradV2,but actually is [%s] and [%s]."
          "The datatype of node[%s]'s two inputs should be equal, but actually is [%d] and [%d], no need to do fusion.",
          elem_node[0]->GetName().c_str(), elem1_node[0]->GetName().c_str(), elem_node[0]->GetType().c_str(),
          elem1_node[0]->GetType().c_str(), elem_node[0]->GetName().c_str(),
          elem_node[0]->GetOpDesc()->GetInputDesc(0).GetDataType(),
          elem_node[0]->GetOpDesc()->GetInputDesc(1).GetDataType());
      return SUCCESS;
    }
  } else {
    if (elem_node[0]->GetType() != "ReluGradV2") {
      FE_LOGI(
          "The optype of node[%s] should be ReluGradV2,"
          "but actually is [%s], no need to do fusion.",
          elem_node[0]->GetName().c_str(), elem_node[0]->GetType().c_str());
      return SUCCESS;
    }
  }
  fusion_nodes = GetMatchedNodes(mapping);
  FE_LOGD("End to do Conv2dBackpropInputElemwiseFusion!");
  return SUCCESS;
}
}  // namespace fe
